# 可以自己import我们平台支持的第三方python模块，比如pandas、numpy等。
import pandas as pd


# 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
def init(context):
    # 选股数
    context.stocknum = 20

    # 升序因子
    context.up = ["market_cap", "pe_ratio", "pb_ratio"]
    # 运行按月定时函数
    scheduler.run_monthly(score_select, tradingday=1)


def score_select(context, bat_dict):
    """打分法选股逻辑"""
    # 选出银子数据, 进行缺失值处理
    q = query(
        fundamentals.eod_derivative_indicator.market_cap,
        fundamentals.eod_derivative_indicator.pe_ratio,
        fundamentals.eod_derivative_indicator.pb_ratio,
        fundamentals.financial_indicator.return_on_invested_capital,
        fundamentals.financial_indicator.inc_revenue,
        fundamentals.financial_indicator.inc_profit_before_tax
    )
    fund = get_fundamentals(q)
    factors_data = fund.T  # 转置
    factors_data = factors_data.dropna()
    # 调试输出
    print(factors_data.head())
    # 定义打分函数, 确定股票池
    select_stocklist(context, factors_data)
    # 定义调仓函数
    reblance(context)
    # 定义打分函数, 确定股票池


def select_stocklist(context, factors_data):
    """打分的具体步骤, 返回股票池"""
    # 循环每个因子去处理
    for name in factors_data.columns:
        # 因子升序的, 进行升序排序
        if name in context.up:
            factor = factors_data.sort_values(by=name)[name]
        else:
            factor = factors_data.sort_values(by=name, ascending=False)[name]
        # 对单个因子进行打分处理
        # 新建一个因子分数列
        factor = pd.DataFrame(factor)
        factor[name + "score"] = 0
        # 进行打分
        # 先求出每组数量, 然后根据数量一次给出分数
        stock_groupnum = len(factors_data) // 10
        for i in range(10):
            # 最后一组单独处理
            if i == 9:
                factor[name + "score"][(i + 1) * stock_groupnum:] = i + 1
            factor[name + "score"][i * stock_groupnum: (i + 1) * stock_groupnum] = i + 1

        # 拼接每个因子的分数到原始的 factors_data 数据中
        factors_data = pd.concat([factors_data, factor[name + "score"]], axis=1)

    # 调试输出
    print(factors_data.head())
    # 求出总分
    sum_score = factors_data[
        ["market_capscore", "pe_ratioscore", "pb_ratioscore", "return_on_invested_capitalscore", "inc_revenuescore",
         "inc_profit_before_taxscore"]].sum(1).sort_values()
    # 调试输出
    print(sum_score.head())
    # 定义股票池
    context.stock_list = sum_score.index[:context.stocknum]
    # 调试输出
    print(context.stock_list)


def reblance(context):
    # ----------------卖出----------------
    # 遍历股票池
    for stock in context.portfolio.positions.keys():
        # 判断是否还在股票池
        if stock not in context.stock_list:
            # 如果不在, 卖出
            order_target_percent(stock, 0)
    # ----------------买入-----------------
    # 买入的百分比
    weight = 1.0 / len(context.stock_list)
    # 遍历股票池
    for stock in context.stock_list:
        # 等比例买入
        order_target_percent(stock, weight)


# before_trading此函数会在每天策略交易开始前被调用，当天只会被调用一次
def before_trading(context):
    pass


# 你选择的证券的数据更新将会触发此段逻辑，例如日或分钟历史数据切片或者是实时数据切片更新
def handle_bar(context, bar_dict):
    # 开始编写你的主要的算法逻辑
    pass


# after_trading函数会在每天交易结束后被调用，当天只会被调用一次
def after_trading(context):
    pass